import pymysql
import json
from tabulate import tabulate
import openai
import re

# 数据库连接信息
connection_config = {
    'host': '10.70.223.112',
    'user': 'dbuser',
    'password': 'Test_12345',
    'database': 'BI_intelligence'
}

def get_database_connection():
    """建立数据库连接并返回连接和游标"""
    connection = pymysql.connect(**connection_config)
    cursor = connection.cursor(pymysql.cursors.DictCursor)
    return connection, cursor

def close_database_connection(connection, cursor):
    """关闭数据库连接和游标"""
    if cursor:
        cursor.close()
    if connection:
        connection.close()

def get_tree_structure():
    """获取树形结构的信息"""
    try:
        connection, cursor = get_database_connection()
        
        # 获取所有Excel文件信息
        cursor.execute("SELECT id, excel_name, excel_describe FROM bi_excel_describe")
        excel_data = cursor.fetchall()
        
        # 获取所有Sheet信息
        cursor.execute("SELECT id, sheet_code, parent_excel_id, sheet_name, sheet_describe FROM bi_sheet_describe")
        sheet_data = cursor.fetchall()
        
        # 构建树形结构
        tree_structure = []
        for excel in excel_data:
            excel_node = {
                "id": excel["id"],
                "name": excel["excel_name"],
                "description": excel["excel_describe"],
                "type": "excel",
                "children": []
            }
            
            # 找到该Excel下的所有Sheet
            for sheet in sheet_data:
                if sheet["parent_excel_id"] == excel["id"]:
                    sheet_node = {
                        "id": sheet["id"],
                        "code": sheet["sheet_code"],
                        "name": sheet["sheet_name"],
                        "description": sheet["sheet_describe"],
                        "type": "sheet",
                        "parent_excel_id": sheet["parent_excel_id"]
                    }
                    excel_node["children"].append(sheet_node)
            
            tree_structure.append(excel_node)
        
        return {
            "excel_data": excel_data,
            "sheet_data": sheet_data,
            "tree_structure": tree_structure
        }
    
    except Exception as e:
        print(f"获取树形结构错误: {e}")
        return None
    
    finally:
        close_database_connection(connection, cursor)

def get_specific_table_info(sheet_code):
    """根据sheet_code获取特定表的信息"""
    try:
        connection, cursor = get_database_connection()
        
        # 构建表名
        table_name = f"{sheet_code}"
        
        # 构建字段映射表名
        mapping_table_name = f"{sheet_code}_describe"
        
        # 检查表是否存在
        cursor.execute("SHOW TABLES LIKE %s", (table_name,))
        table_exists = cursor.fetchone()
        
        if not table_exists:
            return {"error": f"表 {table_name} 不存在"}
        
        # 检查字段映射表是否存在
        cursor.execute("SHOW TABLES LIKE %s", (mapping_table_name,))
        mapping_table_exists = cursor.fetchone()
        
        if not mapping_table_exists:
            print(f"警告: 字段映射表 {mapping_table_name} 不存在，将继续但没有字段映射")
            mapping_data = []
        else:
            # 获取字段映射信息
            cursor.execute(f"SELECT * FROM {mapping_table_name}")
            mapping_data = cursor.fetchall()
        
        # 获取表结构信息
        cursor.execute(f"DESCRIBE {table_name}")
        table_structure = cursor.fetchall()
        
        # 获取表的样本数据
        cursor.execute(f"SELECT * FROM {table_name} LIMIT 5")
        sample_data = cursor.fetchall()
        
        return {
            "table_name": table_name,
            "mapping_data": mapping_data,
            "table_structure": table_structure,
            "sample_data": sample_data
        }
    
    except Exception as e:
        print(f"获取表信息错误: {e}")
        return {"error": str(e)}
    
    finally:
        close_database_connection(connection, cursor)

def create_table_identification_prompt(user_query, tree_info):
    """创建用于识别用户查询涉及哪个表的提示"""
    excel_data = tree_info["excel_data"]
    sheet_data = tree_info["sheet_data"]
    
    # 创建Excel和Sheet的描述信息
    excel_descriptions = ""
    for excel in excel_data:
        excel_descriptions += f"- Excel ID: {excel['id']}, 名称: {excel['excel_name']}, 描述: {excel['excel_describe']}\n"
        
        # 添加该Excel下的所有Sheet
        sheets_for_excel = [sheet for sheet in sheet_data if sheet["parent_excel_id"] == excel["id"]]
        for sheet in sheets_for_excel:
            excel_descriptions += f"  - Sheet ID: {sheet['id']}, 编码: {sheet['sheet_code']}, 名称: {sheet['sheet_name']}, 描述: {sheet['sheet_describe']}\n"
    
    prompt = f"""
我有一个BI系统，包含多个Excel文件和每个Excel下的多个Sheet。以下是系统中的Excel和Sheet的结构:

{excel_descriptions}

用户的查询是: "{user_query}"

请分析这个查询最可能与哪个具体的Sheet相关。

请以下面的JSON格式回答：
{{
  "sheet_code": "最可能的sheet_code",
  "confidence": "高/中/低",
  "reasoning": "简短解释为什么选择这个Sheet"
}}

如果实在无法确定，请设置confidence为"无法确定"，并将sheet_code设为null。
只返回JSON格式的回答，不要有其他任何文字。
"""
    return prompt

def create_sql_generation_prompt(user_query, table_info):
    """创建用于生成SQL的提示"""
    if "error" in table_info:
        return f"错误: {table_info['error']}"
    
    table_name = table_info["table_name"]
    mapping_data = table_info["mapping_data"]
    table_structure = table_info["table_structure"]
    sample_data = table_info["sample_data"]
    
    # 创建字段映射字典
    field_mappings = {}
    for item in mapping_data:
        if 'code' in item and 'name' in item:
            field_mappings[item['code']] = item['name']
    
    # 格式化字段映射信息
    mapping_info = "字段映射信息 (字段名 -> 中文含义):\n"
    for code, name in field_mappings.items():
        mapping_info += f"- {code} -> {name}\n"
    
    # 格式化表结构信息
    structure_info = f"{table_name}表结构:\n"
    for field in table_structure:
        field_name = field.get('Field', '')
        chinese_name = field_mappings.get(field_name, '无映射')
        structure_info += f"- 字段名: {field_name}, 中文含义: {chinese_name}, 类型: {field.get('Type', '')}, 可空: {field.get('Null', '')}\n"
    
    # 格式化样本数据
    sample_info = "数据样本:\n"
    if sample_data:
        headers = sample_data[0].keys()
        rows = [[item[col] for col in headers] for item in sample_data]
        sample_table = tabulate(rows, headers=headers, tablefmt="grid")
        sample_info += sample_table
    
    prompt = f"""
我有一个名为{table_name}的数据表，下面是该表的相关信息：

{structure_info}

{mapping_info}

{sample_info}

用户的自然语言查询是: "{user_query}"

请将用户的自然语言查询转换为针对{table_name}表的SQL查询语句。生成的SQL应该正确反映用户的查询意图。

请严格只返回SQL语句，不需要解释。
"""
    return prompt

def get_response_from_llm(prompt, is_json=False):
    """向大模型发送请求并获取回应"""
    print("\n准备发送给大模型的Prompt:\n")
    print(prompt)
    print("\n====================================================\n")
    
    # 使用本地大模型API
    client = openai.OpenAI(api_key="none", base_url="http://10.70.223.193:8010/v1")
    
    response = client.chat.completions.create(
        model="GLM4",
        messages=[
            {"role": "system", "content": "你是一个专业的数据分析助手。"},
            {"role": "user", "content": prompt},
        ],
        temperature=0.01,
        stream=False,
        max_tokens=4000
    )
    
    answer = response.choices[0].message.content
    
    # 如果需要解析JSON响应
    if is_json:
        # 尝试提取JSON部分
        json_match = re.search(r'({[\s\S]*})', answer)
        if json_match:
            try:
                return json.loads(json_match.group(1))
            except json.JSONDecodeError:
                print("JSON解析失败，返回原始响应")
                return {"error": "JSON解析失败", "raw_response": answer}
        else:
            return {"error": "未找到JSON格式响应", "raw_response": answer}
    else:
        # 清理SQL语句中的markdown格式
        cleaned_answer = answer.replace("```sql\n", "").replace("```sql", "").replace("\n```", "").replace("```", "")
        return cleaned_answer

def execute_sql(sql_query, table_name=None):
    """执行SQL查询并返回结果"""
    try:
        connection, cursor = get_database_connection()
        
        print(f"执行SQL: {sql_query}")
        cursor.execute(sql_query)
        results = cursor.fetchall()
        
        if results:
            headers = results[0].keys()
            rows = [[item[col] for col in headers] for item in results]
            result_table = tabulate(rows, headers=headers, tablefmt="grid")
            return f"查询结果:\n{result_table}"
        else:
            return "查询没有返回任何结果"
    
    except Exception as e:
        return f"SQL执行错误: {e}"
    
    finally:
        close_database_connection(connection, cursor)

def main():
    print("正在获取数据库树形结构信息...")
    tree_info = get_tree_structure()
    
    if not tree_info:
        print("获取树形结构信息失败，请检查连接配置")
        return
    
    print("\n树形结构信息获取成功！")
    print(f"系统中有 {len(tree_info['excel_data'])} 个Excel文件")
    print(f"系统中有 {len(tree_info['sheet_data'])} 个Sheet")
    
    while True:
        print("\n" + "="*60)
        user_query = input("请输入你的自然语言查询 (输入'exit'退出): ")
        
        if user_query.lower() == 'exit':
            print("程序已退出")
            break
        
        # 第一步：确定查询涉及哪个表
        print("\n步骤1: 确定查询涉及哪个表...")
        table_id_prompt = create_table_identification_prompt(user_query, tree_info)
        table_identification = get_response_from_llm(table_id_prompt, is_json=True)
        
        if "error" in table_identification:
            print(f"表识别失败: {table_identification['error']}")
            print(f"原始回应: {table_identification.get('raw_response', '无')}")
            continue
        
        sheet_code = table_identification.get("sheet_code")
        confidence = table_identification.get("confidence", "无法确定")
        reasoning = table_identification.get("reasoning", "无解释")
        
        if not sheet_code or confidence == "无法确定":
            print("无法确定要查询的表，请提供更明确的查询")
            continue
        
        print(f"\n确定查询表: {sheet_code}")
        print(f"置信度: {confidence}")
        print(f"原因: {reasoning}")
        
        # 获取特定表的信息
        print(f"\n步骤2: 获取'{sheet_code}'表的信息...")
        table_info = get_specific_table_info(sheet_code)
        
        if "error" in table_info:
            print(f"获取表信息失败: {table_info['error']}")
            continue
        
        # 第三步：生成SQL查询
        print("\n步骤3: 生成SQL查询...")
        sql_prompt = create_sql_generation_prompt(user_query, table_info)
        sql_query = get_response_from_llm(sql_prompt)
        
        print("\n生成的SQL查询:")
        print(sql_query)
        
        # 询问是否执行查询
        execute = input("\n是否执行此SQL查询? (y/n): ").lower() == 'y'
        if execute:
            result = execute_sql(sql_query, table_info["table_name"])
            print(result)

if __name__ == "__main__":
    main()